# Copyright (c) 2024 Huawei Technologies Co., Ltd.
#
# openMind is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#
#          http://license.coscl.org.cn/MulanPSL2
#
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.
import os
import random
import stat
import time
import unittest
from functools import wraps
from typing import Optional, TypeVar, Callable
import uuid

from openmind_hub.plugins.openmind.om_api import OmApi, create_repo, delete_repo, whoami
from openmind_hub.plugins.openmind.utils._error import AccessWhiteListError


TOKEN = os.environ.get("HUB_TOKEN")

if not TOKEN:
    raise Exception("Please set HUB_TOKEN environment variable.")

ORGANIZATION = "modeltesting"

T = TypeVar("T")

TWENTY_MB = 20 * 1024 * 1024


def random_repo_name(repo_name: Optional[str] = None, prefix: str = "ut_testing") -> str:
    """
    Return a readable pseudo-unique repository name for tests.

    Example:
    ```py
    >>> random_repo_name()
    ut_testing-2fe93f-16599646671840

    >>> random_repo_name("my-space", prefix='space')
    space-my-space-16599481979701
    """
    if repo_name is None:
        repo_name = uuid.uuid4().hex[:6]
    ts = int(time.time() * 10e3)
    return f"{prefix}-{repo_name}-{ts}"


def use_tmp_repo(owner: str = None, repo_type: str = "model") -> Callable[[T], T]:
    """
    Test decorator to create a repo for the test and properly delete it afterward.

    Example:
    ```py
    from openmind.om_api import RepoUrl
    from tests.functional.testing_constants import use_tmp_repo

    class TestOmModel(OmApiCommonTest):
        @use_tmp_repo()
        def test_model_info(self, repo_url: RepoUrl):
            (...)
    ```
    """
    owner = owner or whoami(token=TOKEN).get("username")

    def _inner_use_tmp_repo(test_fn: T) -> T:
        @wraps(test_fn)
        def _inner(*args, **kwargs):
            self = args[0]
            if not isinstance(self, unittest.TestCase):
                raise ValueError("self must be unittest.TestCase")

            repo_url = create_repo(
                repo_id=f"{owner}/{random_repo_name()}", repo_type=repo_type, private=True, token=TOKEN
            )
            try:
                return test_fn(*args, **kwargs, repo_url=repo_url)
            finally:
                delete_repo(repo_id=repo_url.repo_id, repo_type=repo_type, token=TOKEN)

        return _inner

    return _inner_use_tmp_repo


# 创建指定大小的文件
def create_large_file(file_path, size_in_bytes):
    flags = os.O_WRONLY | os.O_CREAT
    modes = stat.S_IWUSR | stat.S_IRUSR
    with os.fdopen(os.open(file_path, flags, modes), "wb") as file:
        bytes_written = 0
        while bytes_written < size_in_bytes:
            # 生成随机数据
            data = bytearray(random.getrandbits(8) for _ in range(1024))
            # 写入文件
            file.write(data)
            bytes_written += len(data)


def is_network_ok():
    api = OmApi(token=TOKEN)
    try:
        api.whoami()
        return True
    except AccessWhiteListError:
        return False
